# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collections
import logging
import re
from dataclasses import dataclass

import paddle
import paddle.distributed as dist
from paddle import pir
from paddle.autograd.backward_utils import ValueDict
from paddle.base.framework import EagerParamBase, pir_op_role_guard
from paddle.base.log_helper import get_logger
from paddle.distributed.fleet.meta_optimizers.common import OpRole
from paddle.distributed.passes.pass_base import PassContext, new_pass
from paddle.distributed.passes.pass_utils import infer_chunk_id

from .mix_to_dist_pass import dist_skip_op_list
from .process_group import get_process_group
from .reshard_funcs.base_reshard_func import (
    choose_reshard_func,
    copy_dist_attr_with_new_member,
    copy_op_attr_with_new_member,
    copy_process_mesh_with_new_member,
)
from .reshard_funcs.reshard_func_register import register_reshard_funcs
from .utils import (
    _complete_op_dist_attr,
    fuse_param_func,
    get_pp_stage_by_pp_degree,
    get_pp_stage_by_process_mesh,
    get_sub_process_mesh_by_program,
    partition_skip_op_list,
    update_pylayer_output,
)

_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
)

register_reshard_funcs()

amp_ops = ["pd_op.check_finite_and_unscale_", "pd_op.update_loss_scaling_"]


def reshard_single_value(program, op, operand, attr):
    prev_var = operand.source()
    if prev_var.is_dist() and prev_var.dist_attr() != attr:
        operand_attr = attr.as_tensor_dist_attr()
        paddle.pir.set_insertion_point(op)
        with pir_op_role_guard(op.op_role):
            # fold reshard
            if prev_var.get_defining_op().name() == 'dist_op.reshard':
                prev_reshard = prev_var.get_defining_op()
                prev_reshard_input = prev_reshard.operand_source(0)
                prev_reshard_result = prev_reshard.result(0)
                # skil global to sub mesh reshard
                if (
                    prev_reshard_input.dist_attr().process_mesh.ndim
                    == prev_reshard_result.dist_attr().process_mesh.ndim
                ):
                    if prev_reshard_input.dist_attr() == operand_attr:
                        return prev_reshard_input
                    reshard_var = paddle._C_ops.reshard_v2(
                        prev_reshard_input, operand_attr
                    )
                    return reshard_var
            # insert reshard
            reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
            return reshard_var
    return prev_var


def reshard_combine_value(program, op, operand, attr):
    prev_var = operand.source()

    assert prev_var.get_defining_op().name() == 'builtin.combine', (
        f"TensorList must be defined by builtin.combine op, but is {prev_var.get_defining_op().name()}."
    )

    combine_op = prev_var.get_defining_op()
    array_attr = attr.as_array_attr()

    assert len(combine_op.operands()) == len(array_attr), (
        "The number of combine op operands and the number of dist array_attr are not equal in op"
    )

    reshard_vars = []
    for inner_operand, inner_attr in zip(combine_op.operands(), array_attr):
        reshard_vars.append(
            reshard_single_value(program, op, inner_operand, inner_attr)
        )

    paddle.pir.set_insertion_point(op)
    with pir_op_role_guard(op.op_role):
        combine_value = paddle._C_ops.builtin_combine(reshard_vars)
    return combine_value


def apply_partition_pass(program, block=None):
    if block is None:
        block = program.global_block()
    for op in block.ops:
        for sub_block in op.blocks():
            apply_partition_pass(program, block=sub_block)

        if op.dist_attr is None:
            continue
        if op.name() in partition_skip_op_list:
            continue

        assert len(op.operands()) == len(op.dist_attr.operands()), (
            f"The number of operands and the number of op_dist_attr's operands are not equal in op: {op}"
        )
        assert len(op.results()) == len(op.dist_attr.results()), (
            f"The number of results and the number of op_dist_attr's results are not equal in op: {op}"
        )

        # deal with inplace value
        for out_idx, in_idx in paddle.core.pir.get_op_inplace_info(op).items():
            ref_op_role = op.op_role

            operand = op.operand(in_idx)
            operand_attr = op.dist_attr.operand(in_idx)
            prev_var = operand.source()
            if (
                not prev_var.is_dist()
                or operand_attr == prev_var.dist_attr()
                or not prev_var.persistable
            ):
                continue

            assert not prev_var.is_combine(), (
                f"The current partition pass not support inplace value of {op} is tensor list."
            )

            operand_attr = operand_attr.as_tensor_dist_attr()

            # reshard input
            paddle.pir.set_insertion_point(op)
            with pir_op_role_guard(ref_op_role):
                reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
                operand.set_source(reshard_var)

            result = op.result(out_idx)
            result_attr = op.dist_attr.result(out_idx).as_tensor_dist_attr()
            assert operand_attr == result_attr, (
                f"For inplace value, The operend dist attr should be equal to result dist attr , please check your infer_spmd func of {op}"
            )

            # reshard output
            paddle.pir.set_insertion_point_after(op)
            old_dist_attr = result.dist_attr()
            result.update_dist_attr(result_attr)

            with pir_op_role_guard(ref_op_role):
                prev_op = prev_var.get_defining_op()

                # reshard output to assign out input
                reshard_var_1 = paddle._C_ops.reshard_v2(
                    result, prev_var.dist_attr()
                )
                assign_out = paddle._C_ops.assign_out_(reshard_var_1, prev_var)
                assign_out.get_defining_op().dist_attr = (
                    copy_op_attr_with_new_member(
                        assign_out.get_defining_op().dist_attr,
                        new_chunk_id=op.dist_attr.chunk_id,
                    )
                )

            if old_dist_attr == result.dist_attr():
                continue

            reshard_var_2 = reshard_var_1
            if old_dist_attr != reshard_var_1.dist_attr():
                with pir_op_role_guard(ref_op_role):
                    reshard_var_2 = paddle._C_ops.reshard_v2(
                        result, old_dist_attr
                    )

            result.replace_all_uses_with(reshard_var_1)
            reshard_var_1.get_defining_op().operand(0).set_source(result)
            reshard_var_2.get_defining_op().operand(0).set_source(result)

        for operand, attr in zip(op.operands(), op.dist_attr.operands()):
            if not attr:
                continue
            prev_var = operand.source()
            if prev_var.is_combine():
                operand.set_source(
                    reshard_combine_value(program, op, operand, attr)
                )
            else:
                operand.set_source(
                    reshard_single_value(program, op, operand, attr)
                )
            prev_op = prev_var.get_defining_op()
            if prev_op and prev_op.num_results() == 1 and prev_var.use_empty():
                prev_op.erase()

        for var, attr in zip(op.results(), op.dist_attr.results()):
            if (
                attr
                and var.initialized()
                and var.is_dist()
                and var.dist_attr() != attr
            ):
                paddle.pir.set_insertion_point_after(op)
                old_dist_attr = var.dist_attr()
                var.update_dist_attr(attr.as_tensor_dist_attr())

                # insert reshard
                with pir_op_role_guard(op.op_role):
                    reshard_var = paddle._C_ops.reshard_v2(var, old_dist_attr)
                    var.replace_all_uses_with(reshard_var)
                    reshard_var.get_defining_op().operand(0).set_source(var)
                    var.get_defining_op().set_bool_attr(
                        "replace_all_uses_with_reshard_var", True
                    )


class ReshardPasses:
    @staticmethod
    def decompose_reshard_pass(dist_program):
        # split composed reshard op into atomic reshard ops, which would increase the opportunity of reshard Re-Use in following fold_reshard_pass.
        del_ops = []
        for op in dist_program.global_block().ops:
            if op.name() != 'dist_op.reshard':
                continue
            input = op.operand_source(0)
            result = op.result(0)

            # split the reshard compose p2p and collective into one p2p reshard and one collective reshard.
            # avoid global to sub mesh case
            if (
                (
                    input.dist_attr().process_mesh
                    != result.dist_attr().process_mesh
                )
                and input.dist_attr().process_mesh.ndim
                == result.dist_attr().process_mesh.ndim
            ):
                if (
                    input.dist_attr().placements
                    != result.dist_attr().placements
                ):
                    ref_op_role = op.op_role
                    with pir_op_role_guard(ref_op_role):
                        intermediate_dist_attr = copy_dist_attr_with_new_member(
                            input.dist_attr(),
                            new_process_mesh=result.dist_attr().process_mesh,
                        )
                        intermediate_dist_type = (
                            paddle.base.libpaddle.pir.cvt_to_dist_type(
                                input.type(), intermediate_dist_attr
                            )
                        )
                        paddle.pir.set_insertion_point(op)
                        intermediate_var = paddle._C_ops.reshard_v2(
                            input, intermediate_dist_attr
                        )
                        new_reshard_result = paddle._C_ops.reshard_v2(
                            intermediate_var, result.dist_attr()
                        )
                        result.replace_all_uses_with(new_reshard_result)
                        del_ops.append(op)

        for op in del_ops:
            _logger.info(f"[Reshard Pass] atomic composed reshard op: {op!s}")
            op.erase()

    @staticmethod
    def fold_reshard_pass(dist_program):
        del_ops = []
        value_dict = ValueDict()
        for op in dist_program.global_block().ops:
            if op.name() != 'dist_op.reshard':
                continue
            input = op.operand_source(0)
            result = op.result(0)
            if input.type() == result.type():
                result.replace_all_uses_with(input)
                del_ops.append(op)
                continue
            if input not in value_dict:
                value_dict[input] = [(result.type(), result)]
                continue
            no_find = True
            for type, val in value_dict[input]:
                if type == result.type():
                    result.replace_all_uses_with(val)
                    del_ops.append(op)
                    no_find = False
                    break
            if no_find:
                value_dict[input].append((result.type(), result))
        for op in del_ops:
            op.erase()

    @staticmethod
    def reshard_op_pass(dist_program, global_params_grads=None, block=None):
        if block is None:
            block = dist_program.global_block()
        for op in block.ops:
            for sub_block in op.blocks():
                ReshardPasses.reshard_op_pass(dist_program, block=sub_block)
            if op.name() == 'dist_op.reshard':
                var = op.operand_source(0)

                op_dist_attr = op.dist_attr
                src_dist_attr = op_dist_attr.operand(0).as_tensor_dist_attr()
                dst_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()

                assert (
                    not var.initialized() or var.dist_attr() == src_dist_attr
                ), (
                    f"The dist_attr of reshard op's input and operand should be equal, but got {var.dist_attr()} and {src_dist_attr}"
                )

                if src_dist_attr == dst_dist_attr:
                    op.result(0).replace_all_uses_with(var)
                    if global_params_grads is not None:
                        for idx, (p, g) in enumerate(global_params_grads):
                            if g is not None and g.is_same(op.result(0)):
                                global_params_grads[idx] = (p, var)
                    op.erase()
                    continue

                paddle.pir.set_insertion_point(op)
                ref_op_role = op.op_role

                all_to_all_dim = (
                    dist.auto_parallel.moe_utils._specific_alltoall_dim(
                        var,
                        dst_dist_attr.process_mesh,
                        dst_dist_attr.placements_attr,
                    )
                )

                if all_to_all_dim is not None:
                    with pir_op_role_guard(ref_op_role):
                        out_value = (
                            dist.auto_parallel.moe_utils._pir_nd_mesh_all2all(
                                op.operand_source(0),
                                op.result(0).type(),
                                dst_dist_attr.process_mesh,
                                dst_dist_attr.placements_attr,
                                all_to_all_dim,
                            )
                        )
                else:
                    reshard_func = choose_reshard_func(
                        src_dist_attr, dst_dist_attr
                    )
                    assert reshard_func is not None, (
                        f'There is no reshard function that matches src_dist_attr: {src_dist_attr} and dst_dist_attr: {dst_dist_attr}, {var.get_defining_op()}'
                    )

                    with pir_op_role_guard(ref_op_role):
                        out_value = reshard_func.reshard(
                            src_dist_attr,
                            dst_dist_attr,
                            op.operand_source(0),
                            op.result(0).type(),
                        )

                if out_value is not None:
                    op.result(0).replace_all_uses_with(out_value)

                if op.result(0).use_empty():
                    if global_params_grads is not None:
                        for idx, (p, g) in enumerate(global_params_grads):
                            if g is not None and g.is_same(op.result(0)):
                                global_params_grads[idx] = (
                                    (p, out_value)
                                    if out_value is not None
                                    else (p, var)
                                )
                    op.erase()

    @staticmethod
    def apply_reshard_pass(dist_program, global_params_grads=None):
        ReshardPasses.decompose_reshard_pass(dist_program)
        ReshardPasses.fold_reshard_pass(dist_program)
        ReshardPasses.reshard_op_pass(dist_program, global_params_grads)


# Replace the specific MoE-related dist op with the
# executable op in the dense program. In expert parallelism
# of the MoE model, the process mesh of each expert is
# different. Two specific apis are used to transform the
# input tensor's global process mesh to the experts' local
# process meshes, which will add two dist ops in the program.
# The following two functions are used to replace the two
# dist ops with the executable share_data_ ops.
def replace_moe_sub_mesh_tensors(op):
    cur_rank = paddle.distributed.get_rank()
    in_value = op.operand_source(0)
    out_value = None
    out_idx = -1
    for idx, val in enumerate(op.results()):
        val_mesh = val.dist_attr().process_mesh
        if cur_rank in val_mesh.process_ids:
            assert out_value is None, (
                f'{op} has more than one results on rank {cur_rank}'
            )
            out_value = val
            out_idx = idx

    paddle.pir.set_insertion_point(op)
    local_value = paddle._C_ops.share_data_(in_value)
    local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
        out_value.type(), out_value.dist_attr()
    )
    local_value.set_type(local_value_type)
    out_value.replace_all_uses_with(local_value)

    op_dist_attr = op.dist_attr
    share_data_op = local_value.get_defining_op()
    share_data_op.dist_attr = (
        paddle.base.libpaddle.pir.create_op_dist_attribute(
            op_dist_attr.process_mesh,
            [op_dist_attr.operand(0).as_tensor_dist_attr()],
            [op_dist_attr.result(out_idx).as_tensor_dist_attr()],
        )
    )
    for val in op.results():
        if not val.use_empty():
            update_pylayer_output(val)
    assert all(val.use_empty() for val in op.results())
    op.erase()


def remove_sub_block_unused_inputs(op):
    inputs_size = op.operand_source.num_operands()
    inputs = [op.operand_source(i) for i in range(inputs_size)]
    # remove unused inputs


class RemovePasses:
    @staticmethod
    def remove_other_rank_op_pass(dist_program):
        # pruning op and value not belong to cur rank
        def prune_op(block):
            cur_rank = paddle.distributed.get_rank()
            reverse_block_ops = block.ops[::-1]
            skip_idx = 0
            for idx, op in enumerate(reverse_block_ops):
                if idx < skip_idx:
                    continue
                skip_idx += 1

                if op.name() == "dist_op.moe_sub_mesh_tensors":
                    replace_moe_sub_mesh_tensors(op)
                    continue
                elif op.name() == "dist_op.moe_global_mesh_tensor":
                    replace_moe_global_mesh_tensor(op)
                    continue
                elif op.name() == "cf.tuple_push":
                    stack_create_op = op.operand_source(0).get_defining_op()
                    if stack_create_op.result(2).use_empty():
                        op.erase()
                    continue
                elif op.name() == "cf.yield":
                    continue
                elif op.name() == "pd_op.pylayer":
                    # if the pylayer op is not on the current rank, we should delete it
                    is_cur_rank = False
                    for pylayer_block in list(op.blocks())[::-1]:
                        for sub_block_op in pylayer_block.ops:
                            if (
                                sub_block_op.dist_attr
                                and cur_rank
                                in sub_block_op.dist_attr.process_mesh.process_ids
                            ):
                                is_cur_rank = True
                                break
                    if not is_cur_rank:
                        op.erase()
                        continue
                    for pylayer_block in list(op.blocks())[::-1]:
                        prune_op(pylayer_block)
                    # update pylayer op's inputs
                    op.as_pylayer_op().update_input()
                    continue
                elif op.name() == "dist_op.dtensor_from_local":
                    dtensor_to_local_idx = idx
                    for i in range(idx, len(reverse_block_ops)):
                        if (
                            reverse_block_ops[i].name()
                            == "dist_op.dtensor_to_local"
                        ):
                            dtensor_to_local_idx = i
                            break
                    if (
                        op.dist_attr
                        and cur_rank
                        not in op.dist_attr.process_mesh.process_ids
                    ):
                        for i in range(idx, dtensor_to_local_idx + 1):
                            reverse_block_ops[i].erase()
                    skip_idx = dtensor_to_local_idx + 1
                    continue
                elif op.name() in partition_skip_op_list:
                    can_delete = True
                    for val in op.results():
                        if not val.use_empty():
                            can_delete = False
                    if can_delete:
                        op.erase()
                    continue

                if (
                    op.dist_attr
                    and cur_rank not in op.dist_attr.process_mesh.process_ids
                ):
                    op.erase()
                elif op.name() == "dist_op.reshard":
                    assert op.result(0).use_empty(), (
                        f'There should not have useful dist.reshard op in remove_other_rank_op_pass. but find : {op}'
                    )
                    op.erase()

        prune_op(dist_program.global_block())

        # merge pd.data ops for
        lr_ops = []
        lr_parameters = []
        for op in dist_program.global_block().ops[::-1]:
            if (
                op.name() == 'pd_op.data'
                and "learning_rate" in op.attrs()["name"]
            ):
                lr_ops.append(op)
            if (
                op.name() == 'builtin.parameter'
                and "learning_rate" in op.attrs()["parameter_name"]
            ):
                lr_parameters.append(op)

        if len(lr_ops) > 1:
            lr_value = lr_ops[0].result(0)
            for op in lr_ops[1:]:
                lr = op.result(0)
                lr.replace_all_uses_with(lr_value)
                op.erase()

        if len(lr_parameters) > 1:
            lr_value = lr_parameters[0].result(0)
            for op in lr_parameters[1:]:
                lr = op.result(0)
                lr.replace_all_uses_with(lr_value)
                op.erase()
        for keyword, argument in dist_program.global_block().kwargs().items():
            if argument.use_empty():
                dist_program.global_block().erase_kwarg(keyword)

    @staticmethod
    def remove_no_need_in_startup(startup_program, main_program):
        # 1. vars used in main_program
        main_program_var_names = []
        for key in main_program.global_block().kwargs():
            main_program_var_names.append(key)
        for op in main_program.global_block().ops:
            for var in op.operands_source():
                if var.has_name:
                    main_program_var_names.append(var.name)
            for var in op.results():
                if var.has_name:
                    main_program_var_names.append(var.name)
        # 2. remove var op not used in main_program
        for op in startup_program.global_block().ops:
            for var in op.operands_source():
                if var.has_name and var.name not in main_program_var_names:
                    op.erase()

        # 3. dead code elimination
        pm = pir.PassManager()
        pm.add_pass('dead_code_elimination_pass', {})
        pm.run(startup_program)
        for op in startup_program.global_block().ops:
            if op.name() == "pd_op.coalesce_tensor_":
                if op.result(0).use_empty() and op.result(1).use_empty():
                    op.erase()
        pm.run(startup_program)

    @staticmethod
    def remove_other_rank_input_output_pass(dist_program):
        '''
        Pruning value not belong to cur rank especially used for check_finite_and_unscale
        and update_loss_scaling op in amp.

        For example, w0 on mesh0, w1 on mesh1, before pass, the ops is:
            [w0_g, w1_g], is_finite = check_finite_and_scale([w0_g, w1_g], loss_scaling)
        after pass, on mesh0, the op is:
            [w0_g], is_finite = check_finite_and_scale([w0_g], loss_scaling)

        Note that here we do not set the op_dist_attr, since it is not used afterwards.
        '''
        cur_rank = paddle.distributed.get_rank()
        for op in dist_program.global_block().ops[::-1]:
            if op.name() not in amp_ops:
                continue
            new_vars = []
            combine_op = op.operand_source(0).get_defining_op()
            for inner_operand in (
                op.operand_source(0).get_defining_op().operands()
            ):
                if (
                    cur_rank
                    in inner_operand.source()
                    .dist_attr()
                    .process_mesh.process_ids
                ):
                    new_vars.append(inner_operand.source())
                    continue
            result = op.operand_source(0).get_defining_op().result(0)
            paddle.pir.set_insertion_point_after(combine_op)
            res = paddle._C_ops.builtin_combine(new_vars)
            res.get_defining_op().op_role = op.op_role
            result.replace_all_uses_with(res)
            combine_op.erase()
            # since it is inplace op, set type of output as the same as input
            op.result(0).set_type(res.type())

    @staticmethod
    def remove_other_rank_params_grads_pass(dist_program, dist_params_grads):
        cur_rank_param = []
        cur_rank = paddle.distributed.get_rank()

        for op in dist_program.global_block().ops:
            if op.name() == 'builtin.parameter':
                if cur_rank in op.dist_attr.process_mesh.process_ids:
                    cur_rank_param.append(op.attrs()['parameter_name'])

        need_remove_idx = []
        for idx, (param, grad) in enumerate(dist_params_grads):
            if grad is None:
                continue
            if param.name not in cur_rank_param:
                need_remove_idx.append(idx)

        for idx in need_remove_idx[::-1]:
            dist_params_grads.pop(idx)

    @staticmethod
    def apply_all(
        dist_main_program, dist_startup_program, dist_params_grads=[]
    ):
        RemovePasses.remove_other_rank_input_output_pass(dist_main_program)
        RemovePasses.remove_other_rank_params_grads_pass(
            dist_main_program, dist_params_grads
        )
        _complete_op_dist_attr(dist_main_program)
        RemovePasses.remove_other_rank_op_pass(dist_main_program)
        RemovePasses.remove_no_need_in_startup(
            dist_startup_program, dist_main_program
        )


def replace_moe_global_mesh_tensor(op):
    cur_rank = paddle.distributed.get_rank()
    out_value = op.result(0)
    in_value = None
    in_idx = -1
    for idx, val in enumerate(op.operands_source()):
        val_mesh = val.dist_attr().process_mesh
        if cur_rank not in val_mesh.process_ids:
            continue
        assert in_value is None, (
            f'{op} has more than one inputs on rank {cur_rank}'
        )
        in_value = val
        in_idx = idx

    paddle.pir.set_insertion_point(op)
    local_value = paddle._C_ops.share_data_(in_value)
    # local_value = paddle.assign(in_value)
    local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
        out_value.type(), out_value.dist_attr()
    )
    local_value.set_type(local_value_type)
    out_value.replace_all_uses_with(local_value)

    op_dist_attr = op.dist_attr
    share_data_op = local_value.get_defining_op()
    share_data_op.dist_attr = (
        paddle.base.libpaddle.pir.create_op_dist_attribute(
            op_dist_attr.process_mesh,
            [op_dist_attr.operand(in_idx).as_tensor_dist_attr()],
            [op_dist_attr.result(0).as_tensor_dist_attr()],
        )
    )

    assert all(val.use_empty() for val in op.results())
    op.erase()


# Note: this is the pass in the dense program
comm_ops = [
    "pd_op.all_gather",
    "pd_op.reduce_scatter",
]


def remove_unuseful_comm_op_pass(program):
    for op in program.global_block().ops:
        if op.name() in comm_ops or (
            op.name() == "pd_op.all_reduce"
            and op.int_attr("reduce_type")
            in [dist.ReduceOp.SUM, dist.ReduceOp.MAX]
        ):
            ring_id = op.int_attr("ring_id")
            process_group = get_process_group(ring_id)
            if process_group.nranks == 1:
                op.result(0).replace_all_uses_with(op.operand_source(0))
                op.erase()
        if op.name() == "pd_op.share_data_":
            if op.operand_source(0).has_one_use():
                op.result(0).replace_all_uses_with(op.operand_source(0))
                op.erase()
        if (
            op.name() == "pd_op.cast"
            and op.result(0).dtype == op.operand_source(0).dtype
        ):
            op.result(0).replace_all_uses_with(op.operand_source(0))
            op.erase()


# In sequence_parallel, we need to transpose hidden_states
# from [bs, seq, hidden] to [seq, bs, hidden] to perform
# split and allgather at dim 0.
# The transpose may lead to about 3% performance
# in llama-70B model (tp4pp8).
# We found that, when bs=1, which is the common case in llm
# training, the transpose is equal to reshape.
# So, this pass is to haddle the specific case.
def eliminate_transpose_by_reshape(program):
    for op in program.global_block().ops:
        if (
            op.name() == 'pd_op.transpose'
            or op.name() == 'pd_op.transpose_grad'
        ):
            var = op.operand(0).source()
            rank = len(var.shape)
            perm = op.attrs()['perm']
            perm = [p + rank if p < 0 else p for p in perm]
            # only support transpose dim 0 and dim 1
            expected_perm = [1, 0] + [i + 2 for i in range(rank - 2)]
            if perm == expected_perm and (
                var.shape[0] == 1 or var.shape[1] == 1
            ):
                paddle.pir.set_insertion_point(op)
                transpose_var = op.result(0)
                reshape_var = paddle._C_ops.reshape(var, transpose_var.shape)
                reshape_var.get_defining_op().op_role = op.op_role
                transpose_var.replace_all_uses_with(reshape_var)
                op.erase()
    return program


def complete_op_role(main_program, op_role_scope: list):
    assert len(op_role_scope) == 3 and len(op_role_scope[0]) == 2, (
        "op_role_scope should has the shape[3, 2]"
    )
    forward_op_start = op_role_scope[0][0]
    forward_op_end = op_role_scope[0][1]

    backward_op_start = op_role_scope[1][0]
    backward_op_end = op_role_scope[1][1]

    opt_op_start = op_role_scope[2][0]
    opt_op_end = op_role_scope[2][1]

    global_op_idx = 0
    for blk in main_program.blocks:
        for op in blk.ops:
            if (
                global_op_idx >= forward_op_start
                and global_op_idx < forward_op_end
            ):
                op.op_role = 0
            elif (
                global_op_idx >= backward_op_start
                and global_op_idx < backward_op_end
            ):
                op.op_role = 1
            elif global_op_idx >= opt_op_start and global_op_idx < opt_op_end:
                op.op_role = 2
            else:
                pass
            global_op_idx += 1


def pipeline_pass(dense_main_program, dense_startup_program, pipeline_strategy):
    """
    Pipeline schedule pass for auto parallel. Enables the pipeline parallel scheduling
    strategies like FThenB, 1F1B, VPP, etc.
    """
    import os

    pass_name = pipeline_strategy.schedule_mode
    assert pass_name in [
        "FThenB",
        "1F1B",
        "VPP",
    ], (
        f"pipeline scheduler only support FThenB, 1F1B and VPP now, but receive {pass_name}"
    )

    pass_attr = {}
    pass_attr["num_micro_batches"] = pipeline_strategy.accumulate_steps
    pass_attr["pp_degree"] = pipeline_strategy.pp_degree
    pass_attr["pp_stage"] = get_pp_stage_by_pp_degree(
        pipeline_strategy.pp_degree
    )
    pass_attr["vpp_degree"] = pipeline_strategy.vpp_degree
    pass_attr["split_backward"] = pipeline_strategy.split_backward

    if pass_name == "1F1B":
        # TODO(Ruibiao): Move FLAGS_1f1b_backward_forward_overlap and
        # FLAGS_mp_async_allreduce_in_backward to auto parallel Strategy
        # after these two optimizations are available.
        pass_attr["enable_backward_forward_overlap"] = int(
            os.environ.get("FLAGS_1f1b_backward_forward_overlap", 0)
        )

    pipeline_pass = new_pass("pipeline_scheduler_" + pass_name, pass_attr)
    pass_context = PassContext()
    pipeline_pass.apply(
        dense_main_program,
        dense_startup_program,
        pass_context,
    )
    plan = pass_context.get_attr("plan")
    return plan


def _extract_seg_method(op, seg_method):
    regex = re.compile(seg_method, re.IGNORECASE)
    struct_name = (
        op.attrs()["struct_name"] if op.has_attr("struct_name") else "/"
    )
    m = regex.search(struct_name)
    if not m:
        return None
    return struct_name[m.start(0) :].split("/")[0]


def _get_seg_struct_names(ops, seg_method):
    fwd_start_op_index = 0
    for i, op in enumerate(ops):
        if _extract_seg_method(op, seg_method):
            fwd_start_op_index = i
            break

    total_op_num = len(ops)
    fwd_end_op_index = total_op_num - 1
    for i in reversed(range(total_op_num)):
        if _extract_seg_method(ops[i], seg_method):
            fwd_end_op_index = i
            break

    struct_names = collections.OrderedDict()
    seg_op_mesh = collections.OrderedDict()

    for i in range(fwd_start_op_index, fwd_end_op_index + 1):
        if ops[i].dist_attr is None:
            continue

        struct_name = _extract_seg_method(ops[i], seg_method)
        if struct_name:
            struct_names[struct_name] = 1
            if struct_name in seg_op_mesh:
                assert (
                    seg_op_mesh[struct_name] == ops[i].dist_attr.process_mesh
                ), "The segment's ops should have same process_mesh."

            seg_op_mesh[struct_name] = ops[i].dist_attr.process_mesh
        else:
            if ops[i].name() != "dist_op.reshard":
                raise ValueError(
                    f"The op {ops[i].name()} without seg_method in its struct_name should only be reshard"
                )

    return list(struct_names.keys())


def _analyze_use_custom_mesh(ops, seg_method, pp_degree):
    non_use_custom_mesh = True
    seg_pp_stages = [-1]

    for op in ops:
        if _extract_seg_method(op, seg_method) and op.dist_attr:
            op_mesh = op.dist_attr.process_mesh
            pp_stage = get_pp_stage_by_process_mesh(op_mesh, pp_degree)
            if pp_stage is None:
                continue

            if seg_pp_stages[-1] > pp_stage:
                non_use_custom_mesh = False
                break
            seg_pp_stages.append(pp_stage)

    if not non_use_custom_mesh:
        _logger.info("Cannot Use Auto VPP")
    else:
        _logger.info("Using Auto VPP")

    return non_use_custom_mesh


def _set_process_mesh_and_chunk_id(
    op,
    chunk_process_mesh,
    chunk_id,
    set_input_mesh=False,
    set_output_mesh=False,
):
    def set_var_origin_op_process_mesh(var_origin_op):
        var_origin_op_input_attr = var_origin_op.dist_attr.operands()
        var_origin_op_output_attr = var_origin_op.dist_attr.results()
        var_origin_op_output_attr[0] = var_origin_op_output_attr[
            0
        ].as_tensor_dist_attr()
        var_origin_op_output_attr[0] = (
            paddle.base.libpaddle.pir.create_tensor_dist_attribute(
                chunk_process_mesh,
                var_origin_op_output_attr[0].dims_mapping,
                var_origin_op_output_attr[0].partial_status,
            )
        )

        var_origin_op.dist_attr = (
            paddle.base.libpaddle.pir.create_op_dist_attribute(
                chunk_process_mesh,
                var_origin_op_input_attr,
                var_origin_op_output_attr,
                0,
            )
        )

    def get_var_process_mesh(var):
        var_process_mesh = None
        var_dist_attr = var.dist_attr()

        def get_attr_mesh(var_dist_attr):
            if var_dist_attr:
                if var_dist_attr.as_array_attr():
                    var_array_attr = var_dist_attr.as_array_attr()
                    return var_array_attr[0].as_tensor_dist_attr().process_mesh
                else:
                    return var_dist_attr.process_mesh

        if var_dist_attr:
            var_process_mesh = get_attr_mesh(var_dist_attr)
        elif var.is_combine():
            # NOTE(zhangwl): op var may is vec_type , need get var dist_attr one by one
            var_list = var.type().as_vec_type()
            var_list = var_list.as_list() if var_list is not None else var_list
            var_attr_list = []
            for combine_var in var_list:
                var_dist_attr = combine_var.as_dist_type().dist_attr()
                var_process_mesh = get_attr_mesh(var_dist_attr)
                if var_process_mesh is not None:
                    return var_process_mesh

    def get_var_attr_with_process_mesh(
        var_dist_attr, var_origin_op, process_mesh
    ):
        # Note(luchang): the var generated by builtin.combine will have multiple dist_attr
        if var_dist_attr and var_dist_attr.as_array_attr():
            var_array_attr = var_dist_attr.as_array_attr()
            for i in range(len(var_array_attr)):
                var_dist_attr = var_array_attr[i].as_tensor_dist_attr()
                if op_mesh is not None:
                    if var_dist_attr.process_mesh == op_mesh:
                        var_array_attr[i] = copy_dist_attr_with_new_member(
                            var_dist_attr, new_process_mesh=process_mesh
                        )
                else:
                    var_array_attr[i] = copy_dist_attr_with_new_member(
                        var_dist_attr, new_process_mesh=process_mesh
                    )
            return var_array_attr
        elif var_dist_attr:
            if op_mesh is not None:
                if var_dist_attr.process_mesh == op_mesh:
                    if var_origin_op.name() in [
                        "pd_op.data",
                        "builtin.parameter",
                    ]:
                        set_var_origin_op_process_mesh(var_origin_op)
                    var_attr = copy_dist_attr_with_new_member(
                        var_dist_attr, new_process_mesh=process_mesh
                    )
                    return var_attr
            else:
                var_attr = copy_dist_attr_with_new_member(
                    var_dist_attr, new_process_mesh=process_mesh
                )
                return var_attr
        return var_dist_attr

    def set_var_process_mesh(var, process_mesh):
        var_dist_attr = var.dist_attr()
        var_origin_op = var.get_defining_op()
        if var_dist_attr:
            var_attr = get_var_attr_with_process_mesh(
                var_dist_attr, var_origin_op, process_mesh
            )
            if var_attr is not None:
                var.update_dist_attr(var_attr)
        elif var.is_combine():
            # NOTE(zhangwl): op var may is vec_type , need set var dist_attr one by one
            var_list = var.type().as_vec_type()
            var_list = var_list.as_list() if var_list is not None else var_list
            var_attr_list = []
            for combine_var in var_list:
                var_dist_attr = combine_var.as_dist_type().dist_attr()
                var_attr_list.append(
                    get_var_attr_with_process_mesh(
                        var_dist_attr, var_origin_op, process_mesh
                    )
                )
            var_array_attr = (
                paddle.base.libpaddle.pir.create_array_dist_attribute(
                    var_attr_list
                )
            )
            var.update_dist_attr(var_array_attr)

    def set_attrs_process_mesh(attrs, process_mesh):
        for idx, attr in enumerate(attrs):
            if attr.as_array_attr():
                array_attr = attr.as_array_attr()
                new_array_attr = []
                for i in range(len(array_attr)):
                    tensor_attr = array_attr[i].as_tensor_dist_attr()
                    new_array_attr.append(tensor_attr)
                    if tensor_attr and tensor_attr.process_mesh == op_mesh:
                        new_array_attr[i] = copy_dist_attr_with_new_member(
                            tensor_attr, new_process_mesh=process_mesh
                        )
                attrs[idx] = (
                    paddle.base.libpaddle.pir.create_array_dist_attribute(
                        new_array_attr
                    )
                )
            else:
                tensor_attr = attr.as_tensor_dist_attr()
                if tensor_attr and tensor_attr.process_mesh == op_mesh:
                    attrs[idx] = copy_dist_attr_with_new_member(
                        tensor_attr, new_process_mesh=process_mesh
                    )

    def set_process_mesh(vars, attrs, process_mesh):
        if vars is not None:
            for var in vars:
                set_var_process_mesh(var, process_mesh)
        if attrs is not None:
            set_attrs_process_mesh(attrs, process_mesh)

    op_input_vars = op.operands_source()
    op_output_vars = op.results()

    # NOTE(zhangwl):dist_skip_op do not have op_mesh
    op_mesh = None
    if op.name() in dist_skip_op_list:
        input_var_process_mesh = None
        # NOTE(zhangwl):dist_skip_op output_process_mesh must equal to input_process_mesh
        for var in op_input_vars:
            input_var_process_mesh = get_var_process_mesh(var)
            if input_var_process_mesh is not None:
                break
        if input_var_process_mesh is not None:
            set_process_mesh(op_output_vars, None, input_var_process_mesh)
        return

    op_dist_attr = op.dist_attr
    op_mesh = op_dist_attr.process_mesh
    op_input_attrs = op_dist_attr.operands()
    op_output_attrs = op_dist_attr.results()

    # if op in seq_chunk , vpp need set var and op chunk_process_mesh and chunk_id
    if set_input_mesh:
        set_process_mesh(op_input_vars, op_input_attrs, chunk_process_mesh)
    if set_output_mesh:
        set_process_mesh(op_output_vars, op_output_attrs, chunk_process_mesh)
    if set_input_mesh or set_output_mesh:
        op_mesh = chunk_process_mesh

    op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
        op_mesh,
        op_input_attrs,
        op_output_attrs,
        chunk_id,
    )


def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
    if not pipeline_strategy.enable:
        return

    sub_process_meshes = get_sub_process_mesh_by_program(dist_program)
    pp_degree = pipeline_strategy.pp_degree
    vpp_degree = pipeline_strategy.vpp_degree
    seg_method = pipeline_strategy.vpp_seg_method
    schedule_mode = pipeline_strategy.schedule_mode
    num_chunks = pp_degree * vpp_degree

    if pp_degree < 2 and vpp_degree > 1:
        raise ValueError("VPP schedule mode only can be set in pipeline mode.")
    if vpp_degree > 1 and (not seg_method or schedule_mode != "VPP"):
        raise ValueError(
            "Please set right schedule_mode and vpp_seg_method for VPP."
        )
    if vpp_degree < 2:
        return

    ReshardPasses.fold_reshard_pass(dist_program)
    seg_struct_names = _get_seg_struct_names(
        dist_program.global_block().ops, seg_method
    )
    ops = dist_program.global_block().ops

    # Step2: analysis whether the pp_stage is non-decreasing among segments
    # 1. if non_use_custom_mesh is True, the ops' process_mesh will be changed by vpp strategy
    # 2. if non_use_custom_mesh is False, the ops's process_mesh will not be changed.
    non_use_custom_mesh = _analyze_use_custom_mesh(ops, seg_method, pp_degree)

    # Step3: Get op index boundary, pp_stage, chunk_id, struct_names of each segment
    seg_pp_stages = [i % pp_degree for i in range(num_chunks)]
    seg_chunk_ids = [i // pp_degree for i in range(num_chunks)]
    seg_parts = [0]
    last_struct_name = None
    # stage_ids[i] represents the stage number assigned to the i-th layer.
    stage_ids = []

    for idx, op in enumerate(ops):
        if len(seg_parts) == len(seg_struct_names):
            break
        struct_name = _extract_seg_method(op, seg_method)
        if op.dist_attr is not None and last_struct_name != struct_name:
            pp_stage = get_pp_stage_by_process_mesh(
                op.dist_attr.process_mesh, pp_degree
            )
            if pp_stage is not None:
                stage_ids.append(pp_stage)
            last_struct_name = struct_name
        if struct_name == seg_struct_names[len(seg_parts)]:
            seg_parts.append(idx)
    seg_parts.append(len(ops))

    pp_stage_layer_nums = [0] * pp_degree
    for i in stage_ids:
        pp_stage_layer_nums[i] = pp_stage_layer_nums[i] + 1
    assert all(value >= vpp_degree for value in pp_stage_layer_nums), (
        "The number of layers on each pp_stage must not be less than the vpp_degree in the pp_stage to ensure that each chunk contains at least one layer."
    )

    seg_layer_num = [0] * num_chunks
    for pp_stage in range(
        0, pp_degree
    ):  # Each pp_stage is assigned a number of layers based on user intent.
        pp_stage_layer_num = pp_stage_layer_nums[pp_stage]
        for i in range(0, pp_stage_layer_num):
            # The pp_stage uses a Round robin scheduling algorithm to allocate layers one by one.
            virtual_chunk_id = i % vpp_degree
            real_chunk_id = (virtual_chunk_id) * pp_degree + pp_stage
            seg_layer_num[real_chunk_id] = seg_layer_num[real_chunk_id] + 1

    # Step4: Set the process_mesh of each op
    seg_id = 0
    reshard_ops = []
    previous_seg_parts_end_idx = 0
    for seg_id in range(num_chunks):
        start_idx = seg_parts[previous_seg_parts_end_idx]
        end_idx = seg_parts[previous_seg_parts_end_idx + seg_layer_num[seg_id]]
        pp_stage = seg_pp_stages[seg_id]
        chunk_id = seg_chunk_ids[seg_id]
        struct_name = ",".join(
            seg_struct_names[
                previous_seg_parts_end_idx : previous_seg_parts_end_idx
                + seg_layer_num[seg_id]
            ]
        )
        previous_seg_parts_end_idx = (
            previous_seg_parts_end_idx + seg_layer_num[seg_id]
        )
        process_mesh = sub_process_meshes[pp_stage]

        _logger.info(
            f"stage=[{pp_stage}], chunk_id=[{chunk_id}], layer_name=[{struct_name}]"
        )
        _logger.info(
            f"start op: [{ops[start_idx].name()}], end op: [{ops[end_idx - 1].name()}]"
        )

        skip_idx = 0
        for idx in range(start_idx, end_idx):
            if idx < skip_idx:
                continue

            is_seg_op = _extract_seg_method(ops[idx], seg_method) is not None
            set_mesh = non_use_custom_mesh & is_seg_op

            if ops[idx].name() == "dist_op.reshard":
                reshard_ops.append(ops[idx])
                continue
            elif ops[idx].name() == "dist_op.dtensor_to_local":
                _set_process_mesh_and_chunk_id(
                    ops[idx],
                    process_mesh,
                    chunk_id,
                    set_input_mesh=set_mesh,
                )
                dtensor_from_local_idx = idx + 1
                while (
                    ops[dtensor_from_local_idx].name()
                    != "dist_op.dtensor_from_local"
                ):
                    dtensor_from_local_idx += 1
                for local_op_idx in range(idx + 1, dtensor_from_local_idx):
                    ops[local_op_idx].set_int_attr("chunk_id", chunk_id)

                _set_process_mesh_and_chunk_id(
                    ops[dtensor_from_local_idx],
                    process_mesh,
                    chunk_id,
                    set_output_mesh=set_mesh,
                )
                skip_idx = dtensor_from_local_idx + 1
                continue

            for sub_block in ops[idx].blocks():
                # TODO(luchang): support condition block
                pass

            _set_process_mesh_and_chunk_id(
                ops[idx],
                process_mesh,
                chunk_id,
                set_input_mesh=set_mesh,
                set_output_mesh=set_mesh,
            )
            skip_idx = idx + 1

    # Step5: set right process_mesh for reshard op
    for op in reshard_ops:
        var = op.operand_source(0)

        op_dist_attr = op.dist_attr
        src_dist_attr = op_dist_attr.operand(0).as_tensor_dist_attr()
        dst_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()

        if src_dist_attr == dst_dist_attr:
            op.result(0).replace_all_uses_with(var)
            op.erase()
            continue

        reshard_func = choose_reshard_func(src_dist_attr, dst_dist_attr)
        reshard_func_name = reshard_func.__class__.__name__

        if reshard_func_name == "NdMeshReshardFunction":
            new_process_mesh = var.dist_attr().process_mesh
            new_src_dist_attr = copy_dist_attr_with_new_member(
                src_dist_attr, new_process_mesh=new_process_mesh
            )
            new_dst_dist_attr = copy_dist_attr_with_new_member(
                dst_dist_attr, new_process_mesh=new_process_mesh
            )
            op.dist_attr = copy_op_attr_with_new_member(
                op_dist_attr,
                new_operands=[new_src_dist_attr],
                new_results=[new_dst_dist_attr],
                new_process_mesh=new_process_mesh,
            )
        elif reshard_func_name == "GlobalToSubMeshFunction":
            result_var = op.result(0)
            new_process_mesh = result_var.dist_attr().process_mesh
            new_dst_dist_attr = copy_dist_attr_with_new_member(
                dst_dist_attr, new_process_mesh=new_process_mesh
            )
            op.dist_attr = copy_op_attr_with_new_member(
                op_dist_attr, new_results=[new_dst_dist_attr]
            )
        elif reshard_func_name == "NdMeshReshardFunctionCrossMesh":
            result_var = op.result(0)
            src_process_mesh = var.dist_attr().process_mesh
            dst_process_mesh = result_var.dist_attr().process_mesh
            new_src_dist_attr = copy_dist_attr_with_new_member(
                src_dist_attr, new_process_mesh=src_process_mesh
            )
            new_dst_dist_attr = copy_dist_attr_with_new_member(
                dst_dist_attr, new_process_mesh=dst_process_mesh
            )
            new_process_ids = (
                src_process_mesh.process_ids + dst_process_mesh.process_ids
            )
            new_process_mesh = copy_process_mesh_with_new_member(
                op.dist_attr.process_mesh,
                new_process_ids=new_process_ids,
            )

            op.dist_attr = copy_op_attr_with_new_member(
                op_dist_attr,
                new_operands=[new_src_dist_attr],
                new_results=[new_dst_dist_attr],
                new_process_mesh=new_process_mesh,
            )
        elif reshard_func_name == "SameStatusReshardFunction":
            op.result(0).replace_all_uses_with(var)
            op.erase()
        else:
            raise ValueError(
                f"Unsupported reshard function: {reshard_func_name}, reshard op's dist_attr: {op.dist_attr}"
            )

    # Step6: add reshard op between pipeline chunks
    apply_partition_pass(dist_program)

    for op in startup_program.global_block().ops:
        if op.name() == "builtin.set_parameter":
            param_name = op.str_attr("parameter_name")
            startup_param = op.operand_source(0)
            param = dist_program.get_parameter_value_by_name(param_name)
            if param.dist_attr():
                startup_param.update_dist_attr(param.dist_attr())

    ReshardPasses.fold_reshard_pass(dist_program)


def check_chunk_id(dist_program):
    all_ops = dist_program.global_block().ops
    for idx, op in enumerate(all_ops):
        if op.op_role in [int(OpRole.Forward), int(OpRole.Backward)]:
            if op.name() in dist_skip_op_list:
                continue

            if op.has_attr("chunk_id"):
                # op between dtensor_from_local and dtensor_to_local will
                # be assigned a chunk_id attribute.
                continue
            elif op.name() in [
                "dist_op.dtensor_from_local",
                "dist_op.dtensor_to_local",
            ]:
                # dtensor_from_local and dtensor_to_local ops will be removed after
                # converting the program to a dense program.
                continue
            elif op.dist_attr.chunk_id == -1:
                if op.name() in ["pd_op.data", "builtin.parameter"]:
                    op.dist_attr = copy_op_attr_with_new_member(
                        op.dist_attr, new_chunk_id=0
                    )
                elif op.name() in ["pd_op.full", "pd_op.full_int_array"]:
                    all_used_ops = op.result(0).all_used_ops()
                    for used_op in all_used_ops:
                        if used_op.dist_attr.chunk_id != -1:
                            op.dist_attr = copy_op_attr_with_new_member(
                                op.dist_attr,
                                new_chunk_id=used_op.dist_attr.chunk_id,
                            )
                            break

                else:
                    op_chunk_id = infer_chunk_id(idx, all_ops)
                    op.dist_attr = copy_op_attr_with_new_member(
                        op.dist_attr, new_chunk_id=op_chunk_id
                    )

            if op.dist_attr.chunk_id == -1:
                raise ValueError(
                    f"The chunk_id of op[{op.name()}] is not set. Please check the chunk_id setting."
                )


def check_order(op_list, order):
    pointer = 0
    for item in order:
        if item == "pd_op.add":
            while (
                pointer < len(op_list)
                and op_list[pointer].name() == "pd_op.add"
            ):
                pointer += 1
        else:
            if pointer >= len(op_list) or op_list[pointer].name() != item:
                return False
            pointer += 1
    return True


def is_ffn_pattern(op_list):
    if len(op_list) != 3 and len(op_list) != 5:
        return False
    order = [
        "pd_op.matmul",
        "pd_op.add",
        "pd_op.matmul",
        "pd_op.add",
        "pd_op.swiglu",
    ]
    return check_order(op_list, order)


def is_qkv_pattern(op_list):
    if len(op_list) != 9 and len(op_list) != 12:
        return False
    order = [
        "pd_op.matmul",
        "pd_op.add",
        "pd_op.full_int_array",
        "pd_op.reshape",
        "pd_op.matmul",
        "pd_op.add",
        "pd_op.full_int_array",
        "pd_op.reshape",
        "pd_op.matmul",
        "pd_op.add",
        "pd_op.full_int_array",
        "pd_op.reshape",
    ]
    return check_order(op_list, order)


def get_param_op(program, param_name):
    all_ops = program.global_block().ops
    for i in range(len(all_ops)):
        if (
            all_ops[i].name() == "builtin.set_parameter"
            and all_ops[i].str_attr("parameter_name") == param_name
        ):
            return [all_ops[i], all_ops[i].operand_source(0).get_defining_op()]


@dataclass
class ParamMeta:
    name: str = None
    local_shape: list = None
    local_num_head: int = None
    local_head_dims: int = None


def fuse_attention_ffn_qkv_pass(
    startup_program, main_program, concrete_program, mode="all"
):
    # 0. Prepare the data structure
    pir_param_names = []
    dy_param_names = []
    for i in range(len(concrete_program.parameters[1])):
        dy_param_names.append(concrete_program.parameters[0][i].name)
        pir_param_names.append(concrete_program.parameters[1][i].name)
    fused_pattern_map = {"ffn": [], "qkv": []}
    fusion_map = {"ffn": [], "qkv": []}

    # 1. Traverse main_program, extract all ffn and qkv patterns.
    all_ops = main_program.global_block().ops
    for i in range(len(all_ops)):
        # check ffn pattern
        if mode == "all" or mode == "ffn":
            pat = all_ops[i : i + 3] if i + 3 <= len(all_ops) else all_ops[i:]
            if is_ffn_pattern(pat):
                fused_pattern_map['ffn'].append(pat)
                i = i + 3
                continue
            else:
                pat = (
                    all_ops[i : i + 5] if i + 5 <= len(all_ops) else all_ops[i:]
                )
                if is_ffn_pattern(pat):
                    fused_pattern_map['ffn'].append(pat)
                    i = i + 5
                    continue
        # check qkv pattern
        if mode == "all" or mode == "qkv":
            pat = all_ops[i : i + 9] if i + 9 <= len(all_ops) else all_ops[i:]
            if is_qkv_pattern(pat):
                fused_pattern_map['qkv'].append(pat)
                i = i + 9
                continue
            else:
                pat = (
                    all_ops[i : i + 12]
                    if i + 12 <= len(all_ops)
                    else all_ops[i:]
                )
                if is_qkv_pattern(pat):
                    fused_pattern_map['qkv'].append(pat)
                    i = i + 12
                    continue

    name2pir_param_map = {}

    # 2. Replace all ffn and qkv patterns with fusion patterns, and record the weights after replacement.
    for pat in fused_pattern_map['ffn']:
        if len(pat) == 5:
            mm_gate = pat[0]
            add_gate = pat[1]
            mm_up = pat[2]
            add_up = pat[3]
        else:
            mm_gate = pat[0]
            add_gate = None
            mm_up = pat[1]
            add_up = None

        fusion_w_name = f"fused_{mm_gate.operand_source(1).name}_{mm_up.operand_source(1).name}"
        fusion_map["ffn"].append(
            {
                fusion_w_name: [
                    ParamMeta(mm_gate.operand_source(1).name, None, None, None),
                    ParamMeta(mm_up.operand_source(1).name, None, None, None),
                ]
            }
        )

        fusion_w_dtype = mm_gate.operand_source(1).dtype
        fusion_w_shape = mm_gate.operand_source(1).shape
        fusion_w_shape[-1] += mm_up.operand_source(1).shape[-1]
        fusion_w_process_mesh = mm_gate.operand_source(1).process_mesh
        # Insert fusion parameter
        with paddle.static.program_guard(main_program, startup_program):
            fused_w = paddle.pir.core.create_parameter(
                dtype=fusion_w_dtype,
                shape=fusion_w_shape,
                name=fusion_w_name,
                process_mesh=fusion_w_process_mesh,
                placements=[
                    paddle.distributed.Replicate(),
                    paddle.distributed.Shard(1),
                ],
                initializer=paddle.nn.initializer.Constant(value=0),
            )
            name2pir_param_map[fusion_w_name] = fused_w
        if add_gate is not None and add_up is not None:
            fusion_bias_name = f"fused_{add_gate.operand_source(1).name}_{add_up.operand_source(1).name}"
            fusion_map["ffn"].append(
                {
                    fusion_bias_name: [
                        ParamMeta(
                            add_gate.operand_source(1).name, None, None, None
                        ),
                        ParamMeta(
                            add_up.operand_source(1).name, None, None, None
                        ),
                    ]
                }
            )

            fusion_bias_dtype = add_gate.operand_source(1).dtype
            fusion_bias_shape = add_gate.operand_source(1).shape
            fusion_bias_shape[-1] += add_up.operand_source(1).shape[-1]
            fusion_bias_process_mesh = add_gate.operand_source(1).process_mesh
            # Insert fusion parameter
            with paddle.static.program_guard(main_program, startup_program):
                fused_bias = paddle.pir.core.create_parameter(
                    dtype=fusion_bias_dtype,
                    shape=fusion_bias_shape,
                    name=fusion_bias_name,
                    process_mesh=fusion_bias_process_mesh,
                    placements=[
                        paddle.distributed.Replicate(),
                        paddle.distributed.Shard(0),
                    ],
                    initializer=paddle.nn.initializer.Constant(value=0),
                )
                name2pir_param_map[fusion_bias_name] = fused_bias

        # Insert dst pattern
        paddle.pir.set_insertion_point_after(pat[-1])
        fused_o = paddle.matmul(
            mm_gate.operand_source(0),
            fused_w,
            transpose_x=False,
            transpose_y=False,
        )
        fused_o.get_defining_op().copy_attrs_from(mm_gate)
        if add_gate is not None and add_up is not None:
            fused_o = paddle.add(fused_o, fused_bias)
            fused_o.get_defining_op().copy_attrs_from(add_gate)
        out = paddle.incubate.nn.functional.swiglu(fused_o)
        out.get_defining_op().copy_attrs_from(pat[-1])
        pat[-1].result(0).replace_all_uses_with(out)

    for pat in fused_pattern_map['qkv']:
        if len(pat) == 12:
            mm_q = pat[0]
            add_q = pat[1]
            reshape_q = pat[3]
            mm_k = pat[4]
            add_k = pat[5]
            reshape_k = pat[7]
            mm_v = pat[8]
            add_v = pat[9]
            reshape_v = pat[11]
        else:
            mm_q = pat[0]
            add_q = None
            reshape_q = pat[2]
            mm_k = pat[3]
            add_k = None
            reshape_k = pat[5]
            mm_v = pat[6]
            add_v = None
            reshape_v = pat[8]

        head_dim = [
            reshape_q.result(0).shape[-1],
            reshape_k.result(0).shape[-1],
            reshape_v.result(0).shape[-1],
        ]
        fusion_w_name = f"fused_{mm_q.operand_source(1).name}_{mm_k.operand_source(1).name}_{mm_v.operand_source(1).name}"
        fusion_map["qkv"].append(
            {
                fusion_w_name: [
                    ParamMeta(
                        mm_q.operand_source(1).name,
                        None,
                        None,
                        reshape_q.result(0).shape[-1],
                    ),
                    ParamMeta(
                        mm_k.operand_source(1).name,
                        None,
                        None,
                        reshape_k.result(0).shape[-1],
                    ),
                    ParamMeta(
                        mm_v.operand_source(1).name,
                        None,
                        None,
                        reshape_v.result(0).shape[-1],
                    ),
                ]
            }
        )
        fusion_w_dtype = mm_q.operand_source(1).dtype
        fusion_w_shape = mm_q.operand_source(1).shape
        fusion_w_shape[-1] += (
            mm_k.operand_source(1).shape[-1] + mm_v.operand_source(1).shape[-1]
        )
        fusion_w_process_mesh = mm_q.operand_source(1).process_mesh
        # insert fusion parameter
        with paddle.static.program_guard(main_program, startup_program):
            fused_w = paddle.pir.core.create_parameter(
                dtype=fusion_w_dtype,
                shape=fusion_w_shape,
                name=fusion_w_name,
                process_mesh=fusion_w_process_mesh,
                placements=[
                    paddle.distributed.Replicate(),
                    paddle.distributed.Shard(1),
                ],
                initializer=paddle.nn.initializer.Constant(value=0),
            )
            name2pir_param_map[fusion_w_name] = fused_w
        if add_q is not None and add_k is not None and add_v is not None:
            fusion_bias_name = f"fused_{add_q.operand_source(1).name}_{add_k.operand_source(1).name}_{add_v.operand_source(1).name}"
            fusion_map["qkv"].append(
                {
                    fusion_bias_name: [
                        ParamMeta(
                            add_q.operand_source(1).name,
                            None,
                            None,
                            reshape_q.result(0).shape[-1],
                        ),
                        ParamMeta(
                            add_k.operand_source(1).name,
                            None,
                            None,
                            reshape_k.result(0).shape[-1],
                        ),
                        ParamMeta(
                            add_v.operand_source(1).name,
                            None,
                            None,
                            reshape_v.result(0).shape[-1],
                        ),
                    ]
                }
            )
            fusion_bias_dtype = add_q.operand_source(1).dtype
            fusion_bias_shape = add_q.operand_source(1).shape
            fusion_bias_shape[-1] += (
                add_k.operand_source(1).shape[-1]
                + add_v.operand_source(1).shape[-1]
            )
            fusion_bias_process_mesh = add_q.operand_source(1).process_mesh
            # insert fusion parameter
            with paddle.static.program_guard(main_program, startup_program):
                fused_bias = paddle.pir.core.create_parameter(
                    dtype=fusion_bias_dtype,
                    shape=fusion_bias_shape,
                    name=fusion_bias_name,
                    process_mesh=fusion_bias_process_mesh,
                    placements=[
                        paddle.distributed.Replicate(),
                        paddle.distributed.Shard(0),
                    ],
                    initializer=paddle.nn.initializer.Constant(value=0),
                )
                name2pir_param_map[fusion_bias_name] = fused_bias
        # insert dst pattern
        paddle.pir.set_insertion_point_after(pat[-1])
        fused_o = paddle.matmul(
            mm_q.operand_source(0),
            fused_w,
            transpose_x=False,
            transpose_y=False,
        )
        fused_o.get_defining_op().copy_attrs_from(mm_q)
        if add_q is not None and add_k is not None and add_v is not None:
            fused_o = paddle.add(fused_o, fused_bias)
            fused_o.get_defining_op().copy_attrs_from(add_q)
        out = paddle.reshape(
            fused_o,
            shape=[
                0,
                0,
                reshape_k.result(0).shape[-2],
                int(
                    (
                        reshape_q.result(0).shape[-2]
                        / reshape_k.result(0).shape[-2]
                        + 2
                    )
                    * reshape_q.result(0).shape[-1]
                ),
            ],
        )
        out.get_defining_op().copy_attrs_from(reshape_q)
        reshape_op = out.get_defining_op()
        if reshape_op.has_attr("struct_name"):
            full_int_array_op = reshape_op.operand_source(1).get_defining_op()
            full_int_array_op.set_str_attr(
                "struct_name", reshape_op.attrs()["struct_name"]
            )
        out_q, out_k, out_v = paddle.split(
            out,
            num_or_sections=[
                int(
                    (
                        reshape_q.result(0).shape[-2]
                        / reshape_k.result(0).shape[-2]
                    )
                    * reshape_q.result(0).shape[-1]
                ),
                reshape_k.result(0).shape[-1],
                reshape_v.result(0).shape[-1],
            ],
            axis=-1,
        )
        if reshape_op.has_attr("struct_name"):
            builtin_split_op = out_q.get_defining_op()
            split_op = builtin_split_op.operand_source(0).get_defining_op()
            builtin_split_op.set_str_attr(
                "struct_name", reshape_op.attrs()["struct_name"]
            )
            split_op.set_str_attr(
                "struct_name", reshape_op.attrs()["struct_name"]
            )
            full_int_array_op = split_op.operand_source(1).get_defining_op()
            full_int_array_op.set_str_attr(
                "struct_name", reshape_op.attrs()["struct_name"]
            )
            full_op = split_op.operand_source(2).get_defining_op()
            full_op.set_str_attr(
                "struct_name", reshape_op.attrs()["struct_name"]
            )
        if reshape_q.result(0).shape[-2] != reshape_k.result(0).shape[-2]:
            out_q = paddle.reshape(
                out_q,
                shape=[
                    0,
                    0,
                    reshape_q.result(0).shape[-2],
                    reshape_q.result(0).shape[-1],
                ],
            )
            if builtin_split_op.has_attr("struct_name"):
                reshape_op = out_q.get_defining_op()
                reshape_op.set_str_attr(
                    "struct_name", builtin_split_op.attrs()["struct_name"]
                )
                full_int_array_op = reshape_op.operand_source(
                    1
                ).get_defining_op()
                full_int_array_op.set_str_attr(
                    "struct_name", builtin_split_op.attrs()["struct_name"]
                )

        reshape_q.result(0).replace_all_uses_with(out_q)
        reshape_k.result(0).replace_all_uses_with(out_k)
        reshape_v.result(0).replace_all_uses_with(out_v)

    # 3. Delete src pattern from origin program.
    del_ops = []
    for pat in fused_pattern_map['ffn']:
        for op in reversed(pat):
            del_ops.append(op)
            if op.name() == "pd_op.matmul" or op.name() == "pd_op.add":
                del_ops.append(op.operand_source(1).get_defining_op())
                del_ops.extend(
                    get_param_op(startup_program, op.operand_source(1).name)
                )
    for pat in fused_pattern_map['qkv']:
        for op in reversed(pat):
            del_ops.append(op)
            if op.name() == "pd_op.matmul" or op.name() == "pd_op.add":
                del_ops.append(op.operand_source(1).get_defining_op())
                del_ops.extend(
                    get_param_op(startup_program, op.operand_source(1).name)
                )
    for op in del_ops:
        op.erase()

    # 4. Initialize fused parameters and delete original parameters.
    concated_dy_param_index = []
    # for key, pat_list in fused_name_map.items():
    for key, pat_list in fusion_map.items():
        for pat in pat_list:
            for pir_param, dy_param_list in pat.items():
                # Retrieve the params of ffn and qkv patterns from concrete_program for fusion.
                concated_dy_param_list = []
                for dy_param in dy_param_list:
                    param_index = dy_param_names.index(dy_param.name)
                    concated_dy_param_list.append(
                        concrete_program.parameters[0][param_index]
                    )
                    dy_param.local_shape = (
                        concrete_program.parameters[0][param_index]
                        ._local_value()
                        .shape
                    )
                    if dy_param.local_head_dims is not None:
                        dy_param.local_num_head = (
                            dy_param.local_shape[-1] // dy_param.local_head_dims
                        )
                    concated_dy_param_index.append(param_index)

                dy_param_init = True
                for p in concated_dy_param_list:
                    if not p._local_value()._is_initialized():
                        dy_param_init = False
                        break

                # Fuse params and init pir program fusion params.
                with paddle.base.dygraph.guard():
                    dyparam_dtype = concated_dy_param_list[0].dtype
                    for param in concated_dy_param_list:
                        assert dyparam_dtype == param.dtype, (
                            "The dtypes of dy parameters to be fused are not the same."
                        )

                    dtensor = paddle.zeros(
                        shape=name2pir_param_map[pir_param].shape,
                        dtype=dyparam_dtype,
                    )
                    fused_dy_param = EagerParamBase.from_tensor(dtensor)
                    fused_dy_param = dist.shard_tensor(
                        fused_dy_param,
                        concated_dy_param_list[0].process_mesh,
                        concated_dy_param_list[0].placements,
                    )
                    fused_dy_param.name = pir_param

                    if dy_param_init:
                        if len(dy_param_list) == 3:
                            is_qkv = True
                            num_heads = dy_param_list[0].local_num_head
                            num_key_value_heads = dy_param_list[
                                1
                            ].local_num_head
                        else:
                            is_qkv = False
                            num_heads = None
                            num_key_value_heads = None
                        concated_param = fuse_param_func(
                            [
                                obj._local_value()
                                for obj in concated_dy_param_list
                            ],
                            is_qkv=is_qkv,
                            num_heads=num_heads,
                            num_key_value_heads=num_key_value_heads,
                        )
                        paddle.assign(
                            concated_param, fused_dy_param._local_value()
                        )
                        concated_param._clear()

                    # Pop and release original params from concrete_program
                    for param in concated_dy_param_list:
                        param.get_tensor()._clear()

                concrete_program.parameters[0].append(fused_dy_param)
                concrete_program.parameters[1].append(
                    name2pir_param_map[pir_param]
                )

    concated_dy_param_index.sort(reverse=True)
    for index in concated_dy_param_index:
        concrete_program.parameters[0].pop(index)
        concrete_program.parameters[1].pop(index)

    return fusion_map
